import os
import subprocess
from mutagen.mp4 import MP4, MP4Cover
from pymkv import MKVFile

def set_video_thumbnail(video_path, image_path):
    """设置视频缩略图（封面）"""
    ext = os.path.splitext(video_path)[1].lower()
    
    if ext == '.mp4':
        # 使用mutagen处理MP4
        video = MP4(video_path)
        with open(image_path, 'rb') as f:
            image_data = f.read()
        
        # 判断图片格式
        if image_path.lower().endswith('.png'):
            cover = MP4Cover(image_data, MP4Cover.FORMAT_PNG)
        else:
            cover = MP4Cover(image_data, MP4Cover.FORMAT_JPEG)
        
        video['covr'] = [cover]
        video.save()
        return True
    
    elif ext == '.mkv':
        # 使用pymkv处理MKV
        mkv = MKVFile(video_path)
        mkv.add_attachment(image_path)
        mkv.mux(f"{os.path.splitext(video_path)[0]}_new.mkv")
        os.replace(f"{os.path.splitext(video_path)[0]}_new.mkv", video_path)
        return True
    
    elif ext in ('.avi', '.divx'):
        # 使用FFmpeg重新封装AVI
        output_path = f"{os.path.splitext(video_path)[0]}_new.avi"
        cmd = [
            'ffmpeg',
            '-i', video_path,
            '-i', image_path,
            '-map', '0',          # 复制所有原始流
            '-map', '1',          # 添加封面图片
            '-c', 'copy',         # 复制流不重新编码
            '-disposition:v:1', 'attached_pic',  # 标记为封面
            output_path
        ]
        try:
            subprocess.run(cmd, check=True, capture_output=True)
            os.replace(output_path, video_path)
            return True
        except subprocess.CalledProcessError:
            return False
    
    else:
        raise ValueError(f"Unsupported video format: {ext}")

# 使用示例
if __name__ == "__main__":
    video_file = "E:/JOYSTICK.DRIVER/_real/91CM-103.mp4"   # 替换为你的视频文件
    thumbnail_image = "E:/JOYSTICK.DRIVER/_real/91CM-103.jpg"   # 替换为你的缩略图文件
    
    if set_video_thumbnail(video_file, thumbnail_image):
        print("缩略图设置成功！")
    else:
        print("操作失败，请检查格式和路径")